# coding:UTF-8
# 同时记录IMU和RTK数据，并分别控制电机的线速度和角速度，存储到不同文件
import binascii
import serial
import time
import csv
import os
import platform
from datetime import datetime
import pynmea2

# --- 检测操作系统并设置串口参数 ---
system_type = platform.system()
print(f"检测到操作系统: {system_type}")

# 根据操作系统类型设置默认串口
if system_type == "Windows":
    # Windows下使用COM口
    motor_port = 'COM6'  # 电机控制串口
    imu_port = 'COM7'    # IMU串口
    rtk_port = 'COM5'    # RTK/GPS串口
else:
    # Linux/Unix下使用/dev/ttyUSB*
    motor_port = '/dev/ttyUSB2'  # 电机驱动板串口号
    imu_port = '/dev/ttyUSB0'    # IMU串口号
    rtk_port = '/dev/ttyUSB1'    # RTK GPS串口号

# --- 串口参数定义 ---
# 电机控制串口
motor_params = {
    'port': motor_port,
    'baudrate': 115200,
    'timeout': 0.5
}

# IMU串口
imu_params = {
    'port': imu_port,
    'baudrate': 9600,
    'timeout': 0.5
}

# RTK/GPS串口
rtk_params = {
    'port': rtk_port,
    'baudrate': 115200,
    'timeout': 0.5
}

# --- 全局变量 ---
# IMU数据相关
ACCData = [0.0]*8
GYROData = [0.0]*8
AngleData = [0.0]*8
FrameState = 0  # 帧状态判断
Bytenum = 0  # 当前帧已接收字节数
CheckSum = 0  # 校验和

a = [0.0]*3  # 加速度
w = [0.0]*3  # 角速度
Angle = [0.0]*3  # 角度

# RTK/GPS数据相关
current_heading = None  # 当前航向角
rtk_timestamp = None  # RTK时间戳

# --- 电机控制相关函数 ---
def generate_control_command(linear_speed=0, angular_speed=0, angle_deflection=100):
    """生成同时控制线速度和角速度的指令
    Args:
        linear_speed: 线速度值 (0-10000)
        angular_speed: 角速度值 (0-10000)
        angle_deflection: 角度值 (默认为100)
    Returns:
        str: 控制指令字符串
    """
    angle_hex = f"{angle_deflection & 0xFFFFFFFF:08X}"  # 角度，限制在32位无符号整数范围内
    linear_hex = f"{-linear_speed & 0xFFFFFFFF:08X}"    # 线速度，负值表示反向
    angular_hex = f"{-angular_speed & 0xFFFFFFFF:08X}"  # 角速度，负值表示反向
    
    # 指令格式： "E0 01 00 00 [角度] [线速度] [角速度]"
    command = f"E0010000{angle_hex}{linear_hex}{angular_hex}"
    print(f"控制命令: {command}")
    return command

def send_command(ser, command):
    """发送控制指令到电机
    Args:
        ser: 串口对象
        command: 控制指令字符串
    """
    try:
        # 将十六进制字符串转换为字节
        byte_command = binascii.unhexlify(command)
        ser.write(byte_command)
        print(f"发送命令: {command}")
    except Exception as e:
        print(f"发送指令错误: {e}")

# --- IMU数据处理函数 ---
def DueData(inputdata):
    """处理IMU原始数据
    Args:
        inputdata: IMU发送的原始字节数据
    Returns:
        bool: 是否成功解析出一组完整数据
    """
    global FrameState
    global Bytenum
    global CheckSum
    global a
    global w
    global Angle
    
    for data in inputdata:  # 遍历输入数据
        if FrameState == 0:  # 状态0：寻找帧头
            if data == 0x55 and Bytenum == 0:  # 找到第一个帧头字节0x55
                CheckSum = data
                Bytenum = 1
                continue
            elif data == 0x51 and Bytenum == 1:  # 加速度数据帧
                CheckSum += data
                FrameState = 1
                Bytenum = 2
            elif data == 0x52 and Bytenum == 1:  # 角速度数据帧
                CheckSum += data
                FrameState = 2
                Bytenum = 2
            elif data == 0x53 and Bytenum == 1:  # 角度数据帧
                CheckSum += data
                FrameState = 3
                Bytenum = 2
        elif FrameState == 1:  # 状态1：接收加速度数据
            if Bytenum < 10:  # 读取8个数据字节
                ACCData[Bytenum-2] = data
                CheckSum += data
                Bytenum += 1
            else:
                if data == (CheckSum & 0xff):  # 校验和检查
                    a = get_acc(ACCData)
                CheckSum = 0
                Bytenum = 0
                FrameState = 0
        elif FrameState == 2:  # 状态2：接收角速度数据
            if Bytenum < 10:
                GYROData[Bytenum-2] = data
                CheckSum += data
                Bytenum += 1
            else:
                if data == (CheckSum & 0xff):
                    w = get_gyro(GYROData)
                CheckSum = 0
                Bytenum = 0
                FrameState = 0
        elif FrameState == 3:  # 状态3：接收角度数据
            if Bytenum < 10:
                AngleData[Bytenum-2] = data
                CheckSum += data
                Bytenum += 1
            else:
                if data == (CheckSum & 0xff):
                    Angle = get_angle(AngleData)
                    return True  # 返回True表示成功解析了一组完整数据
                CheckSum = 0
                Bytenum = 0
                FrameState = 0
    return False  # 返回False表示没有解析到完整数据

def get_acc(datahex):
    """计算加速度值
    Args:
        datahex: 加速度的原始数据字节
    Returns:
        tuple: (x, y, z)轴的加速度
    """
    axl = datahex[0]
    axh = datahex[1]
    ayl = datahex[2]
    ayh = datahex[3]
    azl = datahex[4]
    azh = datahex[5]
    k_acc = 16.0
    acc_x = (axh << 8 | axl) / 32768.0 * k_acc
    acc_y = (ayh << 8 | ayl) / 32768.0 * k_acc
    acc_z = (azh << 8 | azl) / 32768.0 * k_acc
    if acc_x >= k_acc:
        acc_x -= 2 * k_acc
    if acc_y >= k_acc:
        acc_y -= 2 * k_acc
    if acc_z >= k_acc:
        acc_z -= 2 * k_acc
    return acc_x, acc_y, acc_z

def get_gyro(datahex):
    """计算角速度值
    Args:
        datahex: 角速度的原始数据字节
    Returns:
        tuple: (x, y, z)轴的角速度
    """
    wxl = datahex[0]
    wxh = datahex[1]
    wyl = datahex[2]
    wyh = datahex[3]
    wzl = datahex[4]
    wzh = datahex[5]
    k_gyro = 2000.0
    gyro_x = (wxh << 8 | wxl) / 32768.0 * k_gyro
    gyro_y = (wyh << 8 | wyl) / 32768.0 * k_gyro
    gyro_z = (wzh << 8 | wzl) / 32768.0 * k_gyro
    if gyro_x >= k_gyro:
        gyro_x -= 2 * k_gyro
    if gyro_y >= k_gyro:
        gyro_y -= 2 * k_gyro
    if gyro_z >= k_gyro:
        gyro_z -= 2 * k_gyro
    return gyro_x, gyro_y, gyro_z

def get_angle(datahex):
    """计算角度值
    Args:
        datahex: 角度的原始数据字节
    Returns:
        tuple: (x, y, z)轴的角度
    """
    rxl = datahex[0]
    rxh = datahex[1]
    ryl = datahex[2]
    ryh = datahex[3]
    rzl = datahex[4]
    rzh = datahex[5]
    k_angle = 180.0
    angle_x = (rxh << 8 | rxl) / 32768.0 * k_angle
    angle_y = (ryh << 8 | ryl) / 32768.0 * k_angle
    angle_z = (rzh << 8 | rzl) / 32768.0 * k_angle
    if angle_x >= k_angle:
        angle_x -= 2 * k_angle
    if angle_y >= k_angle:
        angle_y -= 2 * k_angle
    if angle_z >= k_angle:
        angle_z -= 2 * k_angle
    return angle_x, angle_y, angle_z

# --- RTK数据处理函数 ---
def process_hdt_data(data):
    """处理HDT语句，获取真实航向角
    Args:
        data: NMEA语句字符串
    Returns:
        float: 航向角
    """
    try:
        sentence = pynmea2.parse(data)
        if isinstance(sentence, pynmea2.HDT):
            true_heading = float(sentence.heading)
            return true_heading
    except (pynmea2.nmea.ChecksumError, pynmea2.ParseError):
        # 忽略ChecksumError和ParseError异常
        pass
    except Exception as e:
        # 捕获并打印其他异常
        print(f"处理HDT数据错误: {e}")
    return None

def read_rtk_data(ser_rtk):
    """读取RTK数据，提取航向角
    Args:
        ser_rtk: RTK串口对象
    Returns:
        heading: 航向角
        timestamp: 时间戳
    """
    global current_heading, rtk_timestamp
    
    if ser_rtk.in_waiting > 0:
        data = ser_rtk.readline().decode('utf-8', errors='ignore').rstrip()
        # print(f"原始数据: {data}")  # 调试时取消注释
        heading = process_hdt_data(data)
        if heading is not None:
            current_heading = heading
            rtk_timestamp = datetime.now().strftime('%H:%M:%S.%f')
            return current_heading, rtk_timestamp
    
    return current_heading, rtk_timestamp

# --- 主程序 ---
if __name__ == '__main__':
    try:
        # 允许用户指定串口
        use_default = input(f"是否使用默认串口设置？(y/n): ")
        if use_default.lower() != 'y':
            motor_params['port'] = input(f"请输入电机控制串口 (默认: {motor_port}): ") or motor_port
            imu_params['port'] = input(f"请输入IMU串口 (默认: {imu_port}): ") or imu_port
            rtk_params['port'] = input(f"请输入RTK/GPS串口 (默认: {rtk_port}): ") or rtk_port
        
        # 打开电机串口
        ser_motor = serial.Serial(**motor_params)
        print(f"电机串口已打开: {ser_motor.is_open}, 端口: {motor_params['port']}")
        
        # 打开IMU串口
        ser_imu = serial.Serial(**imu_params)
        print(f"IMU串口已打开: {ser_imu.is_open}, 端口: {imu_params['port']}")
        
        # 打开RTK串口
        ser_rtk = serial.Serial(**rtk_params)
        print(f"RTK串口已打开: {ser_rtk.is_open}, 端口: {rtk_params['port']}")
        
        # 获取电机线速度
        linear_input = input("请输入电机线速度值(0-10000，输入'q'退出): ")
        if linear_input.lower() == 'q':
            print("程序退出")
            exit(0)
            
        # 线速度检查
        try:
            linear_speed = int(linear_input)
            if not 0 <= linear_speed <= 10000:
                print("线速度超出范围，应该在0-10000之间。默认设置为0。")
                linear_speed = 0
        except ValueError:
            print("输入无效，应为0-10000的整数。默认设置为0。")
            linear_speed = 0
        
        # 获取电机角速度
        angular_input = input("请输入电机角速度值(0-10000，输入'q'退出): ")
        if angular_input.lower() == 'q':
            print("程序退出")
            exit(0)
            
        # 角速度检查
        try:
            angular_speed = int(angular_input)
            if not 0 <= angular_speed <= 10000:
                print("角速度超出范围，应该在0-10000之间。默认设置为0。")
                angular_speed = 0
        except ValueError:
            print("输入无效，应为0-10000的整数。默认设置为0。")
            angular_speed = 0
        
        # 获取文件名基础
        filename_base = input("请输入记录标识(如'test1'): ")
            
        # 创建IMU和RTK数据文件
        imu_filename = f"imu_{filename_base}_L{linear_speed}_A{angular_speed}.csv"
        rtk_filename = f"rtk_{filename_base}_L{linear_speed}_A{angular_speed}.csv"
        print(f"IMU数据将保存到: {imu_filename}")
        print(f"RTK数据将保存到: {rtk_filename}")
        
        # 创建文件并写入表头
        imu_file = open(imu_filename, 'w', newline='')
        imu_writer = csv.writer(imu_file)
        imu_writer.writerow([
            'Timestamp', 
            'IMU_Acc_X', 'IMU_Acc_Y', 'IMU_Acc_Z', 
            'IMU_Gyro_X', 'IMU_Gyro_Y', 'IMU_Gyro_Z',
            'IMU_Angle_X', 'IMU_Angle_Y', 'IMU_Angle_Z',
            'Linear_Speed', 'Angular_Speed'
        ])
        
        rtk_file = open(rtk_filename, 'w', newline='')
        rtk_writer = csv.writer(rtk_file)
        rtk_writer.writerow([
            'Timestamp', 
            'RTK_Heading', 
            'RTK_Timestamp',
            'Linear_Speed', 'Angular_Speed'
        ])
        
        # 发送控制命令启动电机
        print(f"启动电机，线速度: {linear_speed}, 角速度: {angular_speed}")
        motor_command = generate_control_command(linear_speed, angular_speed)
        send_command(ser_motor, motor_command)
        time.sleep(1)  # 等待电机响应
        
        # 设置运行时间
        run_time = 10.0  # 10秒
        start_time = time.time()
        imu_data_count = 0
        rtk_data_count = 0
        
        print(f"开始记录数据，将运行{run_time}秒...")
        
        # 主循环
        while time.time() - start_time < run_time:
            current_time = time.time()
            timestamp = current_time - start_time
            
            # 读取IMU数据
            datahex = ser_imu.read(33)
            
            # 处理IMU数据
            if datahex and DueData(datahex):
                # 显示IMU数据 - 修复格式化语句
                print("加速度: %10.3f %10.3f %10.3f" % (a[0], a[1], a[2]))
                print("角速度: %10.3f %10.3f %10.3f" % (w[0], w[1], w[2]))
                print("角度: %10.3f %10.3f %10.3f" % (Angle[0], Angle[1], Angle[2]))
                print("-" * 50)
                
                # 记录IMU数据
                imu_writer.writerow([
                    timestamp,
                    a[0], a[1], a[2],
                    w[0], w[1], w[2],
                    Angle[0], Angle[1], Angle[2],
                    linear_speed, angular_speed
                ])
                imu_file.flush()
                imu_data_count += 1
                
                if imu_data_count % 10 == 0:
                    remaining = max(0, run_time - (current_time - start_time))
                    print(f"已记录 {imu_data_count} 组IMU数据, 剩余时间: {remaining:.1f}秒")
            
            # 读取RTK数据
            heading, rtk_time = read_rtk_data(ser_rtk)
            
            # 如果读取到有效的航向角数据
            if heading is not None and rtk_time is not None:
                # 显示RTK数据
                print(f"RTK航向角: {heading}°, 时间: {rtk_time}")
                
                # 记录RTK数据
                rtk_writer.writerow([
                    timestamp,
                    heading,
                    rtk_time,
                    linear_speed, angular_speed
                ])
                rtk_file.flush()
                rtk_data_count += 1
                
                if rtk_data_count % 10 == 0:
                    remaining = max(0, run_time - (current_time - start_time))
                    print(f"已记录 {rtk_data_count} 组RTK数据, 剩余时间: {remaining:.1f}秒")
            
            # 短暂休眠以减少CPU使用率
            time.sleep(0.01)
        
        # 停止电机
        print("运行时间结束，停止电机")
        stop_command = "E00100000000000000000000"
        send_command(ser_motor, stop_command)
        
        print(f"共记录 {imu_data_count} 组IMU数据")
        print(f"共记录 {rtk_data_count} 组RTK数据")
        
        # 关闭文件
        imu_file.close()
        rtk_file.close()
        
        # 添加等待用户输入，让用户选择继续记录或退出
        while True:
            print("\n请选择操作:")
            print("1. 设置新的线速度和角速度并继续记录")
            print("2. 仅设置新的线速度(保持相同角速度)")
            print("3. 仅设置新的角速度(保持相同线速度)")
            print("q. 退出程序")
            
            choice = input("请输入选择(1/2/3/q): ")
            
            if choice.lower() == 'q':
                break
            elif choice == '1':
                # 获取新的线速度和角速度
                linear_input = input(f"请输入新的线速度值(0-10000，当前值:{linear_speed}): ")
                try:
                    linear_speed = int(linear_input)
                    if not 0 <= linear_speed <= 10000:
                        print("线速度超出范围，应该在0-10000之间。保持当前值。")
                except ValueError:
                    print("输入无效，保持当前线速度值。")
                
                angular_input = input(f"请输入新的角速度值(0-10000，当前值:{angular_speed}): ")
                try:
                    angular_speed = int(angular_input)
                    if not 0 <= angular_speed <= 10000:
                        print("角速度超出范围，应该在0-10000之间。保持当前值。")
                except ValueError:
                    print("输入无效，保持当前角速度值。")
            elif choice == '2':
                # 仅设置新的线速度
                linear_input = input(f"请输入新的线速度值(0-10000，当前值:{linear_speed}): ")
                try:
                    linear_speed = int(linear_input)
                    if not 0 <= linear_speed <= 10000:
                        print("线速度超出范围，应该在0-10000之间。保持当前值。")
                except ValueError:
                    print("输入无效，保持当前线速度值。")
            elif choice == '3':
                # 仅设置新的角速度
                angular_input = input(f"请输入新的角速度值(0-10000，当前值:{angular_speed}): ")
                try:
                    angular_speed = int(angular_input)
                    if not 0 <= angular_speed <= 10000:
                        print("角速度超出范围，应该在0-10000之间。保持当前值。")
                except ValueError:
                    print("输入无效，保持当前角速度值。")
            else:
                print("无效选择，请重新输入。")
                continue
            
            # 获取文件名基础
            filename_base = input("请输入记录标识: ")
                
            # 新文件名
            imu_filename = f"imu_{filename_base}_L{linear_speed}_A{angular_speed}.csv"
            rtk_filename = f"rtk_{filename_base}_L{linear_speed}_A{angular_speed}.csv"
            print(f"IMU数据将保存到: {imu_filename}")
            print(f"RTK数据将保存到: {rtk_filename}")
            
            # 创建新文件并写入表头
            imu_file = open(imu_filename, 'w', newline='')
            imu_writer = csv.writer(imu_file)
            imu_writer.writerow([
                'Timestamp', 
                'IMU_Acc_X', 'IMU_Acc_Y', 'IMU_Acc_Z', 
                'IMU_Gyro_X', 'IMU_Gyro_Y', 'IMU_Gyro_Z',
                'IMU_Angle_X', 'IMU_Angle_Y', 'IMU_Angle_Z',
                'Linear_Speed', 'Angular_Speed'
            ])
            
            rtk_file = open(rtk_filename, 'w', newline='')
            rtk_writer = csv.writer(rtk_file)
            rtk_writer.writerow([
                'Timestamp', 
                'RTK_Heading', 
                'RTK_Timestamp',
                'Linear_Speed', 'Angular_Speed'
            ])
            
            # 发送控制命令启动电机
            print(f"启动电机，线速度: {linear_speed}, 角速度: {angular_speed}")
            motor_command = generate_control_command(linear_speed, angular_speed)
            send_command(ser_motor, motor_command)
            time.sleep(1)  # 等待电机响应
            
            # 重置计数和时间
            start_time = time.time()
            imu_data_count = 0
            rtk_data_count = 0
            
            print(f"开始记录新数据，将运行{run_time}秒...")
            
            # 主循环
            while time.time() - start_time < run_time:
                current_time = time.time()
                timestamp = current_time - start_time
                
                # 读取IMU数据
                datahex = ser_imu.read(33)
                
                # 处理IMU数据
                if datahex and DueData(datahex):
                    # 显示IMU数据 - 修复格式化语句
                    print("加速度: %10.3f %10.3f %10.3f" % (a[0], a[1], a[2]))
                    print("角速度: %10.3f %10.3f %10.3f" % (w[0], w[1], w[2]))
                    print("角度: %10.3f %10.3f %10.3f" % (Angle[0], Angle[1], Angle[2]))
                    print("-" * 50)
                    
                    # 记录IMU数据
                    imu_writer.writerow([
                        timestamp,
                        a[0], a[1], a[2],
                        w[0], w[1], w[2],
                        Angle[0], Angle[1], Angle[2],
                        linear_speed, angular_speed
                    ])
                    imu_file.flush()
                    imu_data_count += 1
                    
                    if imu_data_count % 10 == 0:
                        remaining = max(0, run_time - (current_time - start_time))
                        print(f"已记录 {imu_data_count} 组IMU数据, 剩余时间: {remaining:.1f}秒")
                
                # 读取RTK数据
                heading, rtk_time = read_rtk_data(ser_rtk)
                
                # 如果读取到有效的航向角数据
                if heading is not None and rtk_time is not None:
                    # 显示RTK数据
                    print(f"RTK航向角: {heading}°, 时间: {rtk_time}")
                    
                    # 记录RTK数据
                    rtk_writer.writerow([
                        timestamp,
                        heading,
                        rtk_time,
                        linear_speed, angular_speed
                    ])
                    rtk_file.flush()
                    rtk_data_count += 1
                    
                    if rtk_data_count % 10 == 0:
                        remaining = max(0, run_time - (current_time - start_time))
                        print(f"已记录 {rtk_data_count} 组RTK数据, 剩余时间: {remaining:.1f}秒")
                
                # 短暂休眠以减少CPU使用率
                time.sleep(0.01)
            
            # 停止电机
            print("运行时间结束，停止电机")
            stop_command = "E00100000000000000000000"
            send_command(ser_motor, stop_command)
            
            print(f"共记录 {imu_data_count} 组IMU数据")
            print(f"共记录 {rtk_data_count} 组RTK数据")
            
            # 关闭文件
            imu_file.close()
            rtk_file.close()
    
    except KeyboardInterrupt:
        print("\n程序被用户中断")
        if 'ser_motor' in locals() and ser_motor.is_open:
            print("紧急停止电机")
            stop_command = "E00100000000000000000000"
            send_command(ser_motor, stop_command)
    
    except serial.SerialException as e:
        print(f"串口错误: {e}")
    
    except Exception as e:
        print(f"程序错误: {e}")
    
    finally:
        # 关闭所有资源
        for file_obj in ['imu_file', 'rtk_file']:
            if file_obj in locals() and not locals()[file_obj].closed:
                locals()[file_obj].close()
                
        for ser_obj in ['ser_motor', 'ser_imu', 'ser_rtk']:
            if ser_obj in locals() and locals()[ser_obj].is_open:
                locals()[ser_obj].close()
                print(f"{ser_obj}串口已关闭") 